package com.boulderai.mcp.client;

import com.boulderai.mcp.model.McpError;
import com.boulderai.mcp.model.McpMessage;
import com.boulderai.mcp.model.Resource;
import com.boulderai.mcp.model.Tool;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 基于WebSocket的MCP客户端实现
 */
public class WebSocketMcpClient implements McpClient {
    
    private static final Logger logger = LoggerFactory.getLogger(WebSocketMcpClient.class);
    
    private final ObjectMapper objectMapper;
    private final AtomicLong requestIdCounter;
    private final Map<Object, CompletableFuture<JsonNode>> pendingRequests;
    
    private WebSocketClient webSocketClient;
    private McpMessageHandler messageHandler;
    private volatile boolean connected = false;
    
    public WebSocketMcpClient() {
        this.objectMapper = new ObjectMapper();
        this.requestIdCounter = new AtomicLong(1);
        this.pendingRequests = new ConcurrentHashMap<>();
    }
    
    @Override
    public CompletableFuture<Void> connect(String serverUri) {
        CompletableFuture<Void> future = new CompletableFuture<>();
        
        try {
            URI uri = new URI(serverUri);
            webSocketClient = new WebSocketClient(uri) {
                @Override
                public void onOpen(ServerHandshake handshake) {
                    logger.info("Connected to MCP server: {}", serverUri);
                    connected = true;
                    future.complete(null);
                }
                
                @Override
                public void onMessage(String message) {
                    handleMessage(message);
                }
                
                @Override
                public void onClose(int code, String reason, boolean remote) {
                    logger.info("Connection closed: code={}, reason={}, remote={}", code, reason, remote);
                    connected = false;
                    if (messageHandler != null) {
                        messageHandler.onClose(code, reason);
                    }
                    // 清理所有待处理的请求
                    pendingRequests.values().forEach(f -> 
                        f.completeExceptionally(new RuntimeException("Connection closed")));
                    pendingRequests.clear();
                }
                
                @Override
                public void onError(Exception ex) {
                    logger.error("WebSocket error", ex);
                    if (messageHandler != null) {
                        messageHandler.onError(ex);
                    }
                    if (!future.isDone()) {
                        future.completeExceptionally(ex);
                    }
                }
            };
            
            webSocketClient.connect();
            
        } catch (Exception e) {
            future.completeExceptionally(e);
        }
        
        return future;
    }
    
    @Override
    public CompletableFuture<Void> disconnect() {
        CompletableFuture<Void> future = new CompletableFuture<>();
        
        if (webSocketClient != null && connected) {
            webSocketClient.close();
            connected = false;
            future.complete(null);
        } else {
            future.complete(null);
        }
        
        return future;
    }
    
    @Override
    public boolean isConnected() {
        return connected && webSocketClient != null && webSocketClient.isOpen();
    }
    
    @Override
    public CompletableFuture<JsonNode> initialize(JsonNode clientInfo) {
        return sendRequest("initialize", clientInfo);
    }
    
    @Override
    public CompletableFuture<List<Tool>> listTools() {
        return sendRequest("tools/list", null)
                .thenApply(result -> {
                    try {
                        JsonNode toolsNode = result.get("tools");
                        return objectMapper.convertValue(toolsNode, new TypeReference<List<Tool>>() {});
                    } catch (Exception e) {
                        throw new RuntimeException("Failed to parse tools list", e);
                    }
                });
    }
    
    @Override
    public CompletableFuture<JsonNode> callTool(String toolName, JsonNode arguments) {
        JsonNode params = objectMapper.createObjectNode()
                .put("name", toolName)
                .set("arguments", arguments);
        return sendRequest("tools/call", params);
    }
    
    @Override
    public CompletableFuture<List<Resource>> listResources() {
        return sendRequest("resources/list", null)
                .thenApply(result -> {
                    try {
                        JsonNode resourcesNode = result.get("resources");
                        return objectMapper.convertValue(resourcesNode, new TypeReference<List<Resource>>() {});
                    } catch (Exception e) {
                        throw new RuntimeException("Failed to parse resources list", e);
                    }
                });
    }
    
    @Override
    public CompletableFuture<JsonNode> readResource(String resourceUri) {
        JsonNode params = objectMapper.createObjectNode()
                .put("uri", resourceUri);
        return sendRequest("resources/read", params);
    }
    
    @Override
    public CompletableFuture<JsonNode> sendRequest(String method, JsonNode params) {
        if (!isConnected()) {
            return CompletableFuture.supplyAsync(() -> {
                throw new IllegalStateException("Not connected to server");
            });
        }
        
        Long requestId = requestIdCounter.getAndIncrement();
        McpMessage message = new McpMessage(requestId, method, params);
        CompletableFuture<JsonNode> future = new CompletableFuture<>();
        
        pendingRequests.put(requestId, future);
        
        try {
            String messageJson = objectMapper.writeValueAsString(message);
            webSocketClient.send(messageJson);
            logger.debug("Sent request: {}", messageJson);
        } catch (Exception e) {
            pendingRequests.remove(requestId);
            future.completeExceptionally(e);
        }
        
        return future;
    }
    
    @Override
    public void sendNotification(String method, JsonNode params) {
        if (!isConnected()) {
            logger.warn("Cannot send notification: not connected to server");
            return;
        }
        
        McpMessage message = new McpMessage(null, method, params);
        
        try {
            String messageJson = objectMapper.writeValueAsString(message);
            webSocketClient.send(messageJson);
            logger.debug("Sent notification: {}", messageJson);
        } catch (Exception e) {
            logger.error("Failed to send notification", e);
        }
    }
    
    @Override
    public void setMessageHandler(McpMessageHandler handler) {
        this.messageHandler = handler;
    }
    
    private void handleMessage(String messageText) {
        try {
            logger.debug("Received message: {}", messageText);
            McpMessage message = objectMapper.readValue(messageText, McpMessage.class);
            
            if (message.isResponse()) {
                handleResponse(message);
            } else if (message.isRequest() || message.isNotification()) {
                if (messageHandler != null) {
                    messageHandler.onMessage(message);
                }
            }
            
        } catch (Exception e) {
            logger.error("Failed to handle message: {}", messageText, e);
            if (messageHandler != null) {
                messageHandler.onError(e);
            }
        }
    }
    
    private void handleResponse(McpMessage message) {
        Long requestId = message.getId();

        if (requestId == null) {
            logger.warn("Received response with null request ID");
            return;
        }
        
        CompletableFuture<JsonNode> future = pendingRequests.get(requestId);
        
        if (future == null) {
            logger.warn("Received response for unknown request ID: {}", requestId);
            return;
        }
        
        // 只有在确认future存在时才移除
        pendingRequests.remove(requestId);
        
        try {
            if (message.getError() != null) {
                McpError error = message.getError();
                future.completeExceptionally(new RuntimeException(
                    String.format("MCP Error %d: %s", error.getCode(), error.getMessage())));
            } else {
                future.complete(message.getResult());
            }
        } catch (Exception e) {
            logger.error("Error completing future for request ID: {}", requestId, e);
            future.completeExceptionally(e);
        }
    }
}